package SurvivalProbability;
use strict;
use List::Util qw(first max maxstr min minstr reduce shuffle sum);

=head1 NAME

SurvivalProbability - Stores the survival probability of a set of species.

=head1 DESCRIPTION

This module stores survival probabilities for a set of species in a format
that can be passed to the NAP optimisation functions. These probabilities 
are specified by a set of cost/survival probability points, between which 
the relationship is assumed to be linear.

=begin comment

 Type    : Constructor
 Title   : _new
 Usage   : my $wx = SurvivalProbability->new(@SpeciesNames);
 Function: Initializes a SurvivalProbability object.
 Alias   :
 Returns : A SurvivalProbability object.
 Args    : A list of species names

=end comment

=cut

sub new {
	
	my $class = shift;
    my $self = {};
    bless $self, $class;    

    #If species names were passed we can set these, otherwise we set a blank species
    #list
    if (@_) { 
        $self->set_names(@_)
    } else {
        $self->{'species_names'} = [];
        $self->set_survivalprobabilities_blank; 
    }
    return $self;
}

#A point is valid for addition
use constant SP_POINT_VALID => 1;
#The constant/species point to be added conflicts with a species/constant point of the
#same value
use constant SP_POINT_CONFLICTS => 4;
#An invalid species was specified
use constant SP_BAD_SPECIES => 8;
#An operation completed successfully
use constant SP_OK => 16;
#A point for deletion was not found 
use constant SP_POINT_NOT_FOUND => 32;
#An operation completed unsuccessfully (use only for checking)
use constant SP_ERROR => 239;

=begin comment

 Type    : Mutator
 Title   : set_names
 Usage   : $wx->set_names(@SpeciesNames);
 Function: Sets the species names and resets survival probabilities
 Returns : 
 Args    : A list of species names

=end comment

=cut

sub set_names {
    my ($self, @names) = @_;
    @names = _unique_sorted_names_(@names);
    #Check that the species names are unique, if not we clear the list
    my %seen = ();
    foreach my $item (@names) {
        if ($seen{$item}++) { 
            #Clear the list of species names
            $self->{'species_names'} = [];
            #Reset the survival probabilities
            $self->set_survivalprobabilities_blank; 
            return SP_ERROR; 
        }
    }

    #Store the species names
    $self->{'species_names'} = \@names;

    #Reset the survival probabilities
    $self->set_survivalprobabilities_blank; 
    return SP_OK;
}


=begin comment

 Type    : Mutator
 Title   : set_default_survival
 Usage   : $wx->set_default_survival
 Function: Sets default survival probabilities of (cost, survival) = (0,0) and (1,1) for all species
 Returns : None
 Args    : None

=end comment

=cut

sub set_default_survival {
	my $self = shift;
	$self->add_point(@{$self->{species_names}},0,0);
	$self->add_point(@{$self->{species_names}},1,1);
}

=begin comment

 Type    : Accessor
 Title   : get_names
 Usage   : $wx->get_names
 Function: Gets a list of species names
 Returns : A list of species names
 Args    : None

=end comment

=cut

sub get_names {
    my $self = shift;
 
    return @{$self->{species_names}};
}


=begin comment

 Type    : Internal function
 Title   : _unique_sorted_names_
 Usage   : $wx->_unique_sorted_names_;
 Function: Returns a list of the unique names, sorted alphabetically
 Returns : The sorted unique names
 Args    : The names

=end comment

=cut

sub _unique_sorted_names_ {
    my @names = @_;
    my @newnames;
    my %seen = ();
    foreach (@names) { unless ($seen{$_}++) { push(@newnames, $_) }; }
    return sort(@newnames);
}

=begin comment

 Type    : Mutator
 Title   : set_survivalprobabilities_blank
 Usage   : $wx->set_survivalprobabilities_blank;
 Function: Resets the survival probabilities
 Returns : 
 Args    : 

=end comment

=cut

sub set_survivalprobabilities_blank {
    my $self = shift;

    #these lists contain the points
    $self->{'survival'}->{'cost'} = [];
    $self->{'survival'}->{'probability'} = [];
    $self->{'survival'}->{'species'} = [];

    #species_points keeps track of which points include this species

    #first blank it
    $self->{'species_points'}={};
    #then add an entry for each species
    foreach (@{$self->{'species_names'}}) {
        $self->{'species_points'}->{$_} = [];
    }
}

=begin comment

 Type    : Mutator
 Title   : add_point
 Usage   : $wx->add_point(@species,$cost,$probability);
 Function: Adds a survival probability/cost point to the species list
 Returns : SP flag 
 Args    : The species list conservation cost and the survival probability

=end comment

=cut

sub add_point {
    my $probability = pop;
    my $cost = pop;
    my ($self, @species) = @_;
    @species = _unique_sorted_names_(@species);

    #Check if this point is a valid point to add
    my $validity = $self->is_point_valid(@species,$cost,$probability);
    unless ($validity & SP_POINT_VALID) { return $validity; }

    #Add the point
    push(@{$self->{'survival'}->{'cost'}}, $cost);
    push(@{$self->{'survival'}->{'probability'}}, $probability);
    push(@{$self->{'survival'}->{'species'}}, \@species);

    #For each effected species make a note of the point number
    foreach (@species) {
        push(@{$self->{'species_points'}->{$_}},(scalar @{$self->{'survival'}->{'cost'}})-1);
    }
    return SP_OK;
}


=begin comment

 Type    : Mutator
 Title   : delete_point
 Usage   : $wx->delete_point(@species,$cost,$probability);
 Function: Deletes a survival probability/cost point. NB: The species list must match 
           the species list for an existing point exactly.
 Returns : SP FLAG
 Args    : The species conservation cost and the survival probability

=end comment

=cut

sub delete_point {
    my $probability = pop;
    my $cost = pop;
    my ($self, @species) = @_;
    @species = _unique_sorted_names_(@species);

    #Check each point that has this set of species
    for (my $index = 0; $index < scalar @{$self->{'survival'}->{'cost'}}; $index++) {
        
        my $match = 1;
        #Check if they have the same number of species
        if (scalar @{$self->{'survival'}->{'species'}->[$index]} != scalar @species) { $match = 0; }
        
        #Check if the species names match
        for (my $sppIndex = 0; $sppIndex < scalar @species; $sppIndex++) {
            if ($species[$sppIndex] ne $self->{'survival'}->{'species'}->[$index]->[$sppIndex]) { $match = 0; }
        }
        
        
        #If we have a match
        if ($self->{'survival'}->{'cost'}->[$index] == $cost && $self->{'survival'}->{'probability'}->[$index] == $probability && $match)
        {
            #Delete the point
            splice(@{$self->{'survival'}->{'cost'}},$index,1);
            splice(@{$self->{'survival'}->{'probability'}},$index,1);
            splice(@{$self->{'survival'}->{'species'}},$index,1);

            #Remove references & update reference numbers for later points
            foreach my $species_name (@{$self->{'species_names'}}) {
                my $toremove = undef;
                for (my $ref_index = 0; $ref_index < scalar @{$self->{'species_points'}->{$species_name}}; $ref_index++) {
                    #This reference will be deleted
                    if ($self->{'species_points'}->{$species_name}->[$ref_index] == $index) { $toremove = $ref_index; }
                    #This point needs to be decremented
                    if ($self->{'species_points'}->{$species_name}->[$ref_index] > $index) { 
                        $self->{'species_points'}->{$species_name}->[$ref_index]--;
                    }
                }

                #Remove the reference
                if (defined $toremove) { splice(@{$self->{'species_points'}->{$species_name}},$toremove,1); }
            }
            return SP_OK;
        }
    } 
    return SP_POINT_NOT_FOUND;
}


=begin comment

 Type    : Accessor
 Title   : is_point_valid
 Usage   : $wx->is_point_invalid(@species,$cost,$probability)
 Function: Check whether this point would be a valid addition
 Returns : SP flag indicating validity
 Args    : The species list conservation cost and the survival probability

=end comment

=cut

sub is_point_valid {
    my $probability = pop;
    my $cost = pop;
    my ($self, @species) = @_;
    @species = _unique_sorted_names_(@species);

    my $out = 0;

    #For each species check if it has a point with this cost already and if that species exists
    foreach my $spp_name (@species) {
        if (exists $self->{'species_points'}->{$spp_name}) { 
            foreach my $point (@{$self->{'species_points'}->{$spp_name}}) {
                if ($self->{'survival'}->{'cost'}->[$point] == $cost) { 
                    $out = $out | SP_POINT_CONFLICTS; 
                }
            }
        } else {
            $out = $out | SP_BAD_SPECIES;
        }
    }

    #If no problems have been detected set the point to valid
    unless ($out) { $out = SP_POINT_VALID; };
    return $out;
}

=begin comment

 Type    : Accessor
 Title   : get_survival_probability
 Usage   : $wx->get_survival_probability($species)
 Function: Gets the survival probability corresponding to different expenditures on each species
 Returns : Hash of survival probabilities
 Args    : Hash containing species expenditures (can use -1 for max possible)

=end comment

=cut
sub get_survival_probability{
	my $self = shift;
	my %set = @_;
	my %output;
	
	foreach (@{$self->{species_names}}) {
		my %survival = $self->get_linear_survival($_);
		if (defined $set{$_}) {
			$output{$_} = $survival{$_}->{b}; 			
		} else {
			$output{$_} = $survival{$_}->{a};
		}
		 
	}
	return %output;
}

=begin comment

 Type    : Accessor
 Title   : get_survival_probability_budget_index
 Usage   : $wx->get_survival_probability($budget, %index)
 Function: Gets the survival probability corresponding to an index on the species
 Returns : Hash of survival probabilities
 Args    : The budget, the species index

=end comment

=cut
sub get_survival_probability_budget_index {
	my $self = shift;
	my $budget = shift;
	my %index = @_;
	
	my %index_r;
	
	#Construct a reverse hash -> the keys are index values the values are
	#arrays containing the species
	foreach (keys %index) { push(@{$index_r{$index{$_}}},$_) }

	my %output;

	#Sorted index values	
	my @index_values = sort keys %index_r;
	
	
	while (scalar @index_values && $budget) {
		my $index = pop(@index_values);
		my $total_expenditure = 0;
		my $survival;
		
		#Get the total cost of conserving all species with this index value completely
		foreach (@{$index_r{$index}}) {
			my @temp_array = $self->get_species_max_survival($_,$budget);
			$output{$_} = shift(@temp_array);
			$total_expenditure += $output{$_}; 
		}
		
		#Get the proportion of the species with this index value that we can conserve
		#and rescale the expenditure on each of those species
		my $proportion; 
		if ($total_expenditure > $budget) {
			#The proportion 
			$proportion = $budget/$total_expenditure; 
			$budget = 0;
		} else {
			$proportion = 1;
			$budget -= $total_expenditure;
		}
		#Get survival probabilities corresponding to rescaled expenditures
		foreach (@{$index_r{$index}}) {
			my @temp_array = $self->get_species_max_survival($_,$output{$_}*$proportion);
			$output{$_} = pop(@temp_array);
		} 
		
	}
	return %output;
}

=begin comment

 Type    : Accessor
 Title   : get_species_max_survival
 Usage   : $wx->get_species_max_survival($species, $budget)
 Function: Finds the highest cost/probability point for a species for a given budget 
 Returns : Conservation cost, survival probability
 Args    : The species, the budget
 
=end comment

=cut
sub get_species_max_survival {
	my $self = shift;
	my $species = shift;
	my $budget = shift;

	my %survival_hash;
	
	$survival_hash{0} = 1;
	
    foreach (@{$self->{'species_points'}->{$species}}) {
		#Store the cost/probability point
    	$survival_hash{$self->{'survival'}->{'cost'}->[$_]} = $self->{'survival'}->{'probability'}->[$_];
    }
    #Default PD style setup 
    if ((scalar keys %survival_hash) == 1) { $survival_hash{1}=1;}
	
	my @sorted = sort keys %survival_hash;
	
	#Budget exceeds maximum cost
	if ($budget >= $sorted[-1]) { return ($sorted[-1], $survival_hash{$sorted[-1]}); }
	#Budget equal to zero
	if ($budget == 0) { return (0,$survival_hash{0}); }
	
	#Budget intermediate
	for (my $index = 0; $index < @sorted; $index++) {
		if ($sorted[$index] <= $budget && $sorted[$index+1] > $budget) {
		 	my $out;
		 	$out = $survival_hash{$sorted[$index]} + 
		 		($survival_hash{$sorted[$index+1]}-$survival_hash{$sorted[$index]}) * 
		 		($budget - $sorted[$index])/($sorted[$index+1]-$sorted[$index]);
			return ($budget, $out);		 			
		}
	} 
	return undef;
}


=begin comment

 Type    : Accessor
 Title   : get_species_cost
 Usage   : $wx->get_species_cost($species)
 Function: Gets all the cost values that apply to a species
 Returns : list of cost values or undef if the species name is invalid
 Args    : The species

=end comment

=cut
sub get_species_cost { return get_species_var(@_,'cost'); }

=begin comment

 Type    : Accessor
 Title   : get_species_probability
 Usage   : $wx->get_species_probability($species)
 Function: Gets all the survival values that apply to a species
 Returns : list of survival values or undef if the species name is invalid
 Args    : $species

=end comment

=cut
sub get_species_probability { return get_species_var(@_,'probability'); }

=begin comment

 Type    : Accessor
 Title   : get_species_species
 Usage   : $wx->get_species_species(@species)
 Function: Gets the species that the points returned by get_species_cost/probability
           also effect
 Returns : list of lists of species also effected by the points
 Args    : @species/$species the species/list of species for which to retrieve points

=end comment

=cut
sub get_species_species { return get_species_var(@_,'species'); }

=begin comment

 Type    : Accessor
 Title   : get_species_var
 Usage   : $wx->get_species_var($species,$var)
 Function: Gets all cost or probability values that apply to a list of species
 Returns : list of cost or probability values or undef if the species name is invalid
 Args    : 

=end comment

=cut
sub get_species_var {
    my $var = pop @_;
    my ($self, @species) = @_;

    #Check a valid field is requested
    unless ( $var =~ m/^(?:probability|cost|species)$/i ) { return undef; }

    #Initially all points are possible
    my @common_points = (0..@{$self->{'survival'}->{$var}});

    #Check species existence and find a common set of points
    foreach my $species_name (@species) {
        #Check that this species exists
        unless (exists $self->{'species_points'}->{$species_name}) { return undef; }
        #Take the intersection of points
        my %original;
        map { $original{$_} = 1 } @common_points;
        @common_points = grep { $original{$_} } @{$self->{'species_points'}->{$species_name}};
    }

    #Retrieve the data
    my @field = (); 
    foreach (@common_points) {
        push(@field,$self->{'survival'}->{$var}->[$_]);
    }
    return @field;
}

=begin comment

 Type    : Accessor
 Title   : get_linear_survival
 Usage   : $wx->get_linear_survival  
 Function: Return a linear survival hash ($self->get_linear_survival should be true!)
 Returns : Linear survival hash
 Args    : None

=end comment

=cut
sub get_linear_survival {
	my $self = shift;
	
	my %out;
	foreach my $species ( @{$self->{species_names}} ) {
		#Defaults are:
		#survival probability 0 for unconserved species
		$out{$species}->{'a'} = 0;
		#survival probability 1 for conserved species
		$out{$species}->{'b'} = 1;
		#conservation cost 1
		$out{$species}->{'c'} = 1;
        #Check out each point
        foreach (@{$self->{'species_points'}->{$species}}) {
        	#If this the zero cost (unconserved) point
        	if ($self->{'survival'}->{'cost'}->[$_] == 0) {
        		$out{$species}->{a} = $self->{'survival'}->{'probability'}->[$_];
       		#If this is the conserved point
        	} else {
        		$out{$species}->{b} = $self->{'survival'}->{'probability'}->[$_];
        		$out{$species}->{c} = $self->{'survival'}->{'cost'}->[$_];
        	}
        }
	}
	return %out;
}

=begin comment

 Type    : Accessor
 Title   : get_linear_survival_tree
 Usage   : $wx->get_linear_survival_tree($tree)  
 Function: Sets generic properties for each terminal node in a tree to indicate it's 
           linear survival function
 Returns : None
 Args    : A Bio::Phylo::Forest::Tree

=end comment

=cut
sub get_linear_survival_tree {
	my $self = shift;
	my $tree = shift;

	foreach my $node ( @{$tree->get_terminals} ) {
		#Defaults are:
		#survival probability 0 for unconserved species
		$node->set_generic(a=>0);
		#survival probability 1 for conserved species
		$node->set_generic(b=>1);
		#conservation cost 1
		$node->set_generic(c=>1);
		
		#Species name
		my $species = $node->get_name;
		
        #Check out each point
        foreach (@{$self->{'species_points'}->{$species}}) {
        	#If this the zero cost (unconserved) point
        	if ($self->{'survival'}->{'cost'}->[$_] == 0) {
        		$node->set_generic( a => $self->{'survival'}->{'probability'}->[$_]);
       		#If this is the conserved point
        	} else {
        		$node->set_generic( b => $self->{'survival'}->{'probability'}->[$_]);
        		$node->set_generic( c => $self->{'survival'}->{'cost'}->[$_]);
        	}
        }
	}
	return 1;
}


=begin comment

 Type    : Mutator
 Title   : set_string
 Usage   : $wx->set_string
 Function: Sets the class from a string representation
 Returns : SP_OK or a string describing the error
 Args    : string

=end comment

=cut
sub set_string {
    my ($self,$string) = @_;

    #Separate string into lines
    my @lines = ($string =~ /.+/g);

    #Get list of species
    my @all_species = ($lines[0] =~ /[^ ]+/g);
    unless ($self->set_names(@all_species) && SP_OK) { return SP_ERROR; }

    #done with the first line
    shift @lines;

    #Process each remaining line
    foreach my $line (@lines) {
        #Get the comma separated fields
        my @fields = ($line =~ /[^,]+/g);
        #Check we have the right number of fields
        unless (@fields % 2 == 1) { return SP_ERROR; };

        if (@fields > 0) {
            #Get the list of species
            my @species = ($fields[0] =~ /[^ ]+/g);
            for (my $fieldIndex = 1; $fieldIndex < @fields; $fieldIndex += 2) {
                unless ($self->add_point(@species,$fields[$fieldIndex],$fields[$fieldIndex+1])) {
                    return SP_ERROR;
                }
            }
        }
    }
    return SP_OK;
}

=begin comment

 Type    : Accessor
 Title   : get_string
 Usage   : $wx->get_string
 Function: Gets a string representation of this class
 Returns : string
 Args    : 

=end comment

=cut
sub get_string {

    my $self = shift;

    #Will contain species lists
    my @lists;
    #Will contain the points for each species list
    my @data;

    #Fill @lists and @data by combining points for each set of species. Looping over the points.
    for (my $pointIndex = 0; $pointIndex < @{$self->{'survival'}->{'cost'}}; $pointIndex++) {

        #See if this set of species has already been entered by comparing with all the lists in @lists so far
        my $match = undef;
        for (my $index = 0; $index < @lists; $index++) { 
            if (_lists_equal_($lists[$index],$self->{'survival'}->{'species'}->[$pointIndex])) { $match = $index; } 
        }
        #If we didn't find a match we make room for the new set of species
        unless (defined $match) {
            #Add the set
            push(@lists,$self->{'survival'}->{'species'}->[$pointIndex]);
            #Add an empty string for this set
            push(@data,'');
            #The matching entry is this one!
            $match = @lists - 1;
        }
        
        #The string has "cost, probability, " appended to it
        $data[$match] .= ', '.$self->{'survival'}->{'cost'}->[$pointIndex].', '.$self->{'survival'}->{'probability'}->[$pointIndex];
    }

    #The final output string
    my $string = '';

    #First line contains a list of species
    foreach (@{$self->{'species_names'}}) { $string .= $_.' '; }
    $string = substr($string,0,-1)."\n";
    
    #For each set of species create a line for the final string
    for (my $lineIndex = 0; $lineIndex < @lists; $lineIndex++) {
        #If all species selected we use _ALL_
        if (@{$lists[$lineIndex]}==@{$self->{'species_names'}}) { $string .= '_ALL_ '; }
        #otherwise we produce a space separated list of species 
        else { foreach (@{$lists[$lineIndex]}) { $string .= $_.' '; } }
        #Then add the data points for that set and a newline
        $string = substr($string,0,-1).$data[$lineIndex]."\n";
    }
    return $string;
}

#Comparing if two lists are equal. There must be a smart/existing way 
#in Perl to do this, however I didn't know and didn't want to require
#an additional package like Test::More
sub _lists_equal_ {
    my @lista = @{$_[0]};
    my @listb = @{$_[1]};

    unless (@lista == @listb) { return 0; }

    for (my $index = 0; $index < @lista; $index++) {
        unless ($lista[$index]  eq $listb[$index]) { return 0; }
    }

    return 1;
}


=begin comment

 Type    : Accessor
 Title   : is_linear
 Usage   : $wx->is_linear
 Function: Check whether this survival probability is linear 
           NB: Linear relationships with multiple points return false
 Returns : true/false
 Args    : none

=end comment

=cut
sub is_linear {
    my $self = shift;
    
    #For each species check it's survival points
    foreach (@{$self->{'species_names'}}) {
        #nonzero_count counts the number of points with a cost > 0
        my $nonzero_count = 0;
        #For each point check if it has a cost > 0
        foreach (@{$self->{'species_points'}->{$_}}) {
            if ($self->{'survival'}->{'cost'}->[$_]) { $nonzero_count++ }
        }
        #If more than one non-zero point return false (non-linear relationship)
        if ($nonzero_count > 1) { return 0; }
    }
    return 1;
}


=begin comment

 Type    : Accessor
 Title   : is_constant_cost
 Usage   : $wx->is_constant_cost
 Function: Check whether this survival probability has constant costs
 Returns : true/false
 Args    : none

=end comment

=cut
sub is_constant_cost {
    my $self = shift;

    my $cost = undef;

    #For each species check its costs
    foreach (@{$self->{'species_names'}}) {
        #For each point check if it has a cost > 0
        foreach (@{$self->{'species_points'}->{$_}}) {
            if ($self->{'survival'}->{'cost'}->[$_]) {  
                if (defined $cost) { 
                    unless ($cost == $self->{'survival'}->{'cost'}->[$_]) { return 0; }
                } else {
                    $cost = $self->{'survival'}->{'cost'}->[$_];
                }
            }
        }
    }
    return 1;
}


=begin comment

 Type    : Accessor
 Title   : is_zero_one_prob
 Usage   : $wx->is_zero_one_prob
 Function: Check whether this survival probability goes from zero to one
 Returns : true/false
 Args    : none

=end comment

=cut
sub is_zero_one_prob {
    my $self = shift;

    my $cost = undef;

    #For each species check its costs
    foreach (@{$self->{'species_names'}}) {
        #For each point check if it has a cost > 0
        foreach (@{$self->{'species_points'}->{$_}}) {
            my $prob = $self->{'survival'}->{'probability'}->[$_];
            unless ($prob == 0 || $prob == 1) { return 0; }
        }
    }
    return 1;
}

=begin comment

 Type    : Accessor
 Title   : is_survival_ratio_constant
 Usage   : $wx->is_survival_ratio_constant
 Function: Check that (1-b)/(1-a) is constant for all species
 Returns : true/false
 Args    : none

=end comment

=cut
sub is_survival_ratio_constant {
    my $self = shift;

    my $ratio = undef;

    #For each species check its survival ratio
    foreach (@{$self->{'species_names'}}) {
        my $a = 0;
        my $b = 1;

        #For each point check if it has a cost > 0
        foreach (@{$self->{'species_points'}->{$_}}) {
            if ($self->{'survival'}->{'cost'}->[$_] > 0) {
                $b = $self->{'survival'}->{'probability'}->[$_];
            } else {
                $a = $self->{'survival'}->{'probability'}->[$_];
            }
            my $prob = $self->{'survival'}->{'probability'}->[$_];
        }

        #If this species is guaranteed to survive without conservation then we cannot 
        #use this algorithm... (other than the uninteresting case where all species 
        #are guaranteed to survive)
        if ($a == 1) { return 0; }

        #Check the ratio
        if (defined $ratio) {
            unless (abs($ratio - (1-$b)/(1-$a))<1e-6) { return 0; }
        } else {
            $ratio = (1-$b)/(1-$a);
        }
    }
    return 1;
}
1;